package stateful;

import org.apache.commons.io.Charsets;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;

import java.io.RandomAccessFile;

/**
 * AtLeastOnce Source
 */
public class MyAtLeastOnceSourceDemo {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(new Configuration());

        // 开启checkpoint
        env.enableCheckpointing(3000);

        env.setParallelism(4);

        // 读取数据
        DataStreamSource<String> lines = env.addSource(new AtLeastOnceSource("data"));

        SingleOutputStreamOperator<String> mapped = lines.map(new MapFunction<String, String>() {
            @Override
            public String map(String value) throws Exception {
                if (value.contains("error")) {
                    throw new RuntimeException("数据出错!");
                }
                return value;
            }
        });
        mapped.print();
        // 直接输出
        lines.print();


        env.execute("");
    }

    private static class AtLeastOnceSource extends RichParallelSourceFunction<String> implements CheckpointedFunction {

        // 是否产生数据
        private boolean flag = true;

        // 偏移量
        private Long offset = 0L;

        // 路径
        private String path;

        // 状态
        private ListState<Long> state;

        public AtLeastOnceSource(String path) {
            this.path = path;
        }

        /**
         * 初始化状态 算子第一次被调用时最先调用
         *
         * @param context
         * @throws Exception
         */
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            System.out.println("Invoked init");
            ListStateDescriptor<Long> stateDescriptor = new ListStateDescriptor<>("list-state", Long.class);
            // 初始化状态
            state = context.getOperatorStateStore().getListState(stateDescriptor);
            // 恢复偏移量
            for (Long l : state.get()) {
                offset = l;
            }
        }

        /**
         * 执行initializeState方法后 调用一次
         *
         * @param parameters
         * @throws Exception
         */
        @Override
        public void open(Configuration parameters) throws Exception {
            System.out.println("Invoked Open");
        }

        /**
         * 产生数据方法
         *
         * @param ctx
         * @throws Exception
         */
        @Override
        public void run(SourceContext<String> ctx) throws Exception {
            // 读取数据
            int indexOfThisSubtask = getRuntimeContext().getIndexOfThisSubtask();
            // 创建文件对象
            RandomAccessFile randomAccessFile = new RandomAccessFile(path + "/" + indexOfThisSubtask + ".txt", "r");
            // 从指定位置读取
            randomAccessFile.seek(offset);
            while (flag) {
                String line = randomAccessFile.readLine();
                if (line != null) {
                    line = new String(line.getBytes(Charsets.ISO_8859_1), Charsets.UTF_8);
                    // 上锁
                    synchronized (ctx.getCheckpointLock()) {
                        // 更新偏移量
                        offset = randomAccessFile.getFilePointer();
                        // 输出数据
                        ctx.collect(indexOfThisSubtask + ".txt: " + line);
                    }
                } else {
                    Thread.sleep(1000);
                }
            }
        }

        /**
         * source停止时调用
         */
        @Override
        public void cancel() {
            flag = false;
        }

        /**
         * checkpoint 周期调用 保存分布式快照
         *
         * @param context
         * @throws Exception
         */
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            // 更新快照
            state.clear();
            state.add(offset);
        }
    }
}


